
#include "camera_adapter.h"
#include "muduo/base/Timestamp.h"

#include "json.hpp"
#include <opencv2/opencv.hpp>

#include <fstream>

#include <dirent.h>

namespace perception::camera {

static std::string timestamp_to_string(const muduo::Timestamp& t)
{
    char buf[64] = { 0 };
    time_t seconds = static_cast<time_t>(t.microSecondsSinceEpoch() / 1000 / 1000);
    struct tm tm_time;
    localtime_r(&seconds, &tm_time);

    snprintf(buf, sizeof(buf), "%4d%02d%02d_%02d%02d%02d_%03d", tm_time.tm_year + 1900,
        tm_time.tm_mon + 1, tm_time.tm_mday, tm_time.tm_hour, tm_time.tm_min, tm_time.tm_sec,
        int((t.microSecondsSinceEpoch() / 1000) % 1000));
    return buf;
}

static muduo::Timestamp last_save_time;
static std::string video_prefix;
static const std::string output_dir = "/home/nvidia/workspace/data/train_data/";
static size_t scene_counter = 0;
static size_t frame_counter = 0;
static std::map<std::string, std::shared_ptr<cv::VideoWriter>> gWriters;
static nlohmann::json jstr;

static std::map<ObjectType, std::string> typid2str_tbl = {
    { ObjectType::UNKNOWN, "unknown" },
    { ObjectType::UNKNOWN_MOVABLE, "movable" },
    { ObjectType::UNKNOWN_UNMOVABLE, "unmovable" },
    { ObjectType::CAR, "car" },
    { ObjectType::VAN, "van" },
    { ObjectType::TRUCK, "truck" },
    { ObjectType::BUS, "bus" },
    { ObjectType::CYCLIST, "cyclist" },
    { ObjectType::MOTORCYCLIST, "motorcyclist" },
    { ObjectType::TRICYCLIST, "tricyclist" },
    { ObjectType::PEDESTRIAN, "pedestrian" },
    { ObjectType::TRAFFICCONE, "trafficcone" },
    { ObjectType::SAFETY_TRIANGLE, "triangle" },
    { ObjectType::BARRIER_DELINEATOR, "delineator" },
    { ObjectType::BARRIER_WATER, "water" },
    { ObjectType::ANIMAL, "animal" },
};

cv::Mat paint_attribute_on_cvmat(const cv::Mat& image, const std::vector<ObjectTrackInfoPtr>& objects,
    const std::vector<std::string>& keys)
{
    auto f2d = [](const Point2f& in) -> cv::Point { return cv::Point(in.x, in.y); };
    cv::Mat out = image.clone();
#if 1
    // paint
    for (auto obj : objects) {
        auto& pts8 = obj->detect.pts8;
        // // 底面四边形(俯视图)
        cv::line(out, f2d(pts8.pts8[0]), f2d(pts8.pts8[1]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[1]), f2d(pts8.pts8[2]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[2]), f2d(pts8.pts8[3]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[3]), f2d(pts8.pts8[0]), cv::Scalar(0, 0xFF, 0xFF), 1);

        // // 顶部四边形
        cv::line(out, f2d(pts8.pts8[4]), f2d(pts8.pts8[5]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[5]), f2d(pts8.pts8[6]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[6]), f2d(pts8.pts8[7]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[7]), f2d(pts8.pts8[4]), cv::Scalar(0, 0xFF, 0xFF), 1);

        // // 竖直线
        cv::line(out, f2d(pts8.pts8[0]), f2d(pts8.pts8[4]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[1]), f2d(pts8.pts8[5]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[2]), f2d(pts8.pts8[6]), cv::Scalar(0, 0xFF, 0xFF), 1);
        cv::line(out, f2d(pts8.pts8[3]), f2d(pts8.pts8[7]), cv::Scalar(0, 0xFF, 0xFF), 1);
        // 是否在图片上增加车辆属性

        {
            // trackid + type
            std::stringstream ss;
            ss << obj->tracker.track_id << ":" << typid2str_tbl[ObjectType(obj->detect.type_id)];
            const auto text = ss.str();
            const cv::Size text_size = cv::getTextSize(text, cv::FONT_HERSHEY_COMPLEX, 1, 1, nullptr);
            const auto rect = cv::Rect(obj->detect.box.left_top.x, obj->detect.box.left_top.y - text_size.height, text_size.width, text_size.height);
            cv::rectangle(out, rect, cv::Scalar(0x00, 0x00, 0x00, 255), cv::FILLED);
            cv::putText(out, text, f2d(obj->detect.box.left_top), cv::FONT_HERSHEY_SCRIPT_SIMPLEX, 1.0, cv::Scalar(0, 0xFF, 0xFF), 1);
        }
    }
#endif
    double acc_height = 0;
#if 0
    for (auto key : keys) { // 图片信息
        const cv::Size key_size = cv::getTextSize(key, cv::FONT_HERSHEY_COMPLEX, 0.5, 1, nullptr);
        const auto rect = cv::Rect(0, 0 + acc_height, key_size.width, key_size.height);
        cv::rectangle(out, rect, cv::Scalar(0x00, 0x00, 0x00), cv::FILLED);
        cv::putText(out, key, cv::Point(0, key_size.height + acc_height), cv::FONT_HERSHEY_COMPLEX, 0.5, cv::Scalar(0xFF, 0xFF, 0xFF), 1);
        acc_height += (key_size.height + 2);
    }
#endif
    return out;
}

static bool isDirectory(const char* path)
{
    DIR* pdir = opendir(path);
    if (pdir != NULL) {
        closedir(pdir);
        pdir = NULL;
        return true;
    }
    return false;
}

static void modifyDirPath(std::string& path) // 修改目录路径为X/Y/Z/
{
    if (path.empty()) {
        return;
    }
    for (std::string::iterator iter = path.begin(); iter != path.end(); ++iter) {
        if (*iter == '\\') {
            *iter = '/';
        }
    }
    if (path.at(path.length() - 1) != '/') {
        path += "/";
    }
}

static bool createRecursionDir(const char* dir)
{
    std::string dirs(dir);
    if (dirs.empty()) {
        return true;
    }

    modifyDirPath(dirs);

    std::string::size_type pos = dirs.find('/');
    while (pos != std::string::npos) {
        std::string cur = dirs.substr(0, pos - 0);
        if (cur.length() > 0 && !isDirectory(cur.c_str())) {
            bool ret = (mkdir(cur.c_str(), S_IRWXU | S_IRWXG | S_IRWXO) == 0);
            if (!ret && errno != EEXIST) {
                return false;
            }
        }
        pos = dirs.find('/', pos + 1);
    }

    return true;
}

static void init_writers(cv::Size sz)
{
    for (auto i = 0; i < 7; i++) {
        const auto camera = "camera" + std::to_string(i);
        const auto file = output_dir + "/" + video_prefix + "/" + video_prefix + "_painted_" + camera + ".mp4";
        gWriters[camera] = std::make_shared<cv::VideoWriter>(file, cv::VideoWriter::fourcc('m', 'p', '4', 'v'), 20, sz);
    }
}

void make_train_data(const std::unordered_map<std::string, cv::Mat>& images,
    std::unordered_map<std::string, std::vector<ObjectTrackInfoPtr>>& tracked)
{
    auto now = muduo::Timestamp::now();

    if (muduo::timeDifference(now, last_save_time) > 10) {
        scene_counter++;
        const auto jfilename = "/Scene_V" + video_prefix + "_S" + std::to_string(scene_counter);
        std::fstream jfile = std::fstream(output_dir + "/" + video_prefix + "/" + jfilename + ".json", std::ios::out | std::ios::trunc);
        jfile << jstr.dump(4);
        jstr.clear();
        last_save_time = now;

        video_prefix = timestamp_to_string(muduo::Timestamp::now());
        createRecursionDir(std::string(output_dir + "/" + video_prefix + "/").c_str());
        cv::Size sz(images.begin()->second.cols, images.begin()->second.rows);
        init_writers(sz);
    }

    frame_counter++;
    for (auto image : images) {
        const auto oname = video_prefix + "_" + image.first + "_TS" + std::to_string(now.microSecondsSinceEpoch() / 1000) + "_F" + std::to_string(frame_counter);
        const auto picname = "V" + oname + ".jpg";
        { // 存图片
            cv::imwrite(output_dir + "/" + video_prefix + "/" + picname, image.second);
        }

        { // 存视频
            const auto scene_name = "Scene_V" + video_prefix + "_S" + std::to_string(scene_counter + 1);
            auto save_image = paint_attribute_on_cvmat(image.second, tracked[image.first], { scene_name, picname });
            gWriters[image.first]->write(save_image);
        }

        auto& objects_this_camera = jstr[std::to_string(now.microSecondsSinceEpoch() / 1000)][image.first];
        objects_this_camera["filename"] = picname;

        size_t object_counter = 0;
        for (auto& obj : tracked[image.first]) {
            const auto& object = obj;
            auto& out = objects_this_camera["objects"][object_counter++];

            out.push_back(typid2str_tbl[ObjectType(object->detect.type_id)]);

            out.push_back(object->detect.box.left_top.x);
            out.push_back(object->detect.box.left_top.y);
            out.push_back(object->detect.box.right_bottom.x);
            out.push_back(object->detect.box.right_bottom.y);

            out.push_back(object->detect.pts8.pts8[0].x);
            out.push_back(object->detect.pts8.pts8[0].y);
            out.push_back(object->detect.pts8.pts8[1].x);
            out.push_back(object->detect.pts8.pts8[1].y);
            out.push_back(object->detect.pts8.pts8[2].x);
            out.push_back(object->detect.pts8.pts8[2].y);
            out.push_back(object->detect.pts8.pts8[3].x);
            out.push_back(object->detect.pts8.pts8[3].y);
            out.push_back(object->detect.pts8.pts8[4].x);
            out.push_back(object->detect.pts8.pts8[4].y);
            out.push_back(object->detect.pts8.pts8[5].x);
            out.push_back(object->detect.pts8.pts8[5].y);
            out.push_back(object->detect.pts8.pts8[6].x);
            out.push_back(object->detect.pts8.pts8[6].y);
            out.push_back(object->detect.pts8.pts8[7].x);
            out.push_back(object->detect.pts8.pts8[7].y);

            out.push_back(object->tracker.length);
            out.push_back(object->tracker.width);
            out.push_back(0);

            out.push_back(object->tracker.track_id);
            std::cout << out.dump() << std::endl;
        }
    }
    std::cout << "**************************** read " << video_prefix << " done!  *****************************" << std::endl;
}
}